import numpy as np
from .base_optimizer import BaseOptimizer

# 假设主类原名为ParticleSwarmOptimizer或其它，添加别名
try:
    PSO = ParticleSwarmOptimizer
except NameError:
    pass
try:
    PSO = PSOOptimizer
except NameError:
    pass
# 如果主类本身就是PSO，则无需更改

class PSO(BaseOptimizer):
    def __init__(self, dim, pop_size=50, max_iter=1000, w=0.7, c1=1.5, c2=1.5):
        super().__init__(dim, pop_size, max_iter)
        self.w = w  # 惯性权重
        self.c1 = c1  # 个体学习因子
        self.c2 = c2  # 社会学习因子
        
    def optimize(self, objective_func, bounds=(-5, 5)):
        # 初始化种群
        population = self.initialize_population(bounds)
        velocities = np.zeros((self.pop_size, self.dim))
        
        # 计算适应度
        fitness = np.array([objective_func(ind) for ind in population])
        
        # 初始化个体最优和全局最优
        pbest = population.copy()
        pbest_fitness = fitness.copy()
        gbest_idx = np.argmin(fitness)
        self.best_solution = population[gbest_idx].copy()
        self.best_fitness = fitness[gbest_idx]
        
        for gen in range(self.max_iter):
            for i in range(self.pop_size):
                # 更新速度
                r1, r2 = np.random.rand(2)
                velocities[i] = (self.w * velocities[i] + 
                               self.c1 * r1 * (pbest[i] - population[i]) +
                               self.c2 * r2 * (self.best_solution - population[i]))
                
                # 更新位置
                population[i] += velocities[i]
                population[i] = self.clip_to_bounds(population[i], bounds)
                
                # 评估新解
                f = objective_func(population[i])
                
                # 更新个体最优
                if f < pbest_fitness[i]:
                    pbest[i] = population[i].copy()
                    pbest_fitness[i] = f
                    
                    # 更新全局最优
                    if f < self.best_fitness:
                        self.best_solution = population[i].copy()
                        self.best_fitness = f
                        
        return self.best_solution, self.best_fitness 